import sys


def verify_cuda_cupy_compatibility():
    """验证CUDA和CuPy的兼容性"""
    try:
        import cupy as cp

        print("\n=== CUDA 和 CuPy 版本验证 ===")
        print(f"Python版本: {sys.version}")
        print(f"CuPy版本: {cp.__version__}")

        if cp.cuda.is_available():
            # CUDA信息
            runtime_version = cp.cuda.runtime.runtimeGetVersion()
            cuda_version = f"{runtime_version // 1000}.{(runtime_version % 1000) // 10}"
            print(f"CUDA版本: {cuda_version}")

            # GPU设备信息
            device_count = cp.cuda.runtime.getDeviceCount()
            print(f"\n=== 发现 {device_count} 个GPU设备 ===")

            for device_id in range(device_count):
                device = cp.cuda.runtime.getDeviceProperties(device_id)
                print(f"\nGPU {device_id}:")
                print(f"  设备名称: {device['name'].decode()}")
                print(f"  计算能力: {device['major']}.{device['minor']}")
                print(f"  总内存: {device['totalGlobalMem'] / (1024 ** 3):.2f} GB")

            # 测试GPU功能
            print("\n=== 执行GPU功能测试 ===")
            try:
                x = cp.array([1, 2, 3])
                _ = cp.sqrt(x)  # 使用下划线表示有意忽略的变量
                print("GPU计算测试: 成功")
            except Exception as e:
                print(f"GPU计算测试失败: {str(e)}")

        else:
            print("警告: CUDA不可用")
            print("请检查:")
            print("1. NVIDIA驱动是否正确安装")
            print("2. CUDA工具包是否正确安装")
            print("3. CuPy是否正确安装")

    except ImportError:
        print("错误: 未安装CuPy")
        print("请运行: pip install cupy-cuda[xx]x")
        print("其中[xx]替换为您的CUDA版本")

    except Exception as e:
        print(f"验证过程出错: {str(e)}")


if __name__ == "__main__":
    verify_cuda_cupy_compatibility()
